import sys
sys.path.append('../../')
import BasicLibraries   as BL 
import Make_Population
import global_variables as GV
import os
import regressions
import warnings
from statsmodels.tools.sm_exceptions import ConvergenceWarning
warnings.simplefilter('ignore', ConvergenceWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=RuntimeWarning)

'''
This file contains the three stage estimation approach described in the main text and the Material and Methods
The functions are used in the main_bootstraped.py file

'''

marker_type, line_style, index_dict, abb_index_dict, inverse_dict, STcontrols = GV.global_vars()
MARGINAL=False
ANNOTATOR='Triple'

############################
data = Make_Population.MakeUniverse()
empiricsBS = regressions.Analysis()
ROB = regressions.RobustnessTest()


def STAGE1(compustat, regions, sector, tc, system_thinking, paris2018,paris2019,paris2020, base_dummies, MAX_ProjectionYear, ref, vars_, cdp_t, ref_t, cdp_firms, non_cdp_firms):
    relevance_effect_expl, relevant_explanators_rob  = ROB.step_1_explanators(compustat, regions, sector, tc, system_thinking, paris2018,paris2019,paris2020, base_dummies, MAX_ProjectionYear, ref, vars_, cdp_t, ref_t, cdp_firms, non_cdp_firms)
    BL.pickle.dump([relevance_effect_expl, relevant_explanators_rob ], open('Relevant_explanators_rob'+ANNOTATOR+'.pckl', 'wb'))
    robust_explanators = relevant_explanators_rob.rename(index = inverse_dict).index.values.tolist()
    
    return robust_explanators


def STAGE2(compustat, regions, sector, tc, system_thinking, paris2018,paris2019,paris2020, base_dummies, MAX_ProjectionYear, ref, robust_explanators, vars_, cdp_t, ref_t, cdp_firms, non_cdp_firms, MARGINAL=False):
    relevance_effect_med, relevant_mediators_rob  = ROB.step_2_mediators(compustat, regions, sector, tc, system_thinking, paris2018,paris2019,paris2020, base_dummies, MAX_ProjectionYear, ref, robust_explanators, vars_, cdp_t, ref_t, cdp_firms, non_cdp_firms, MARGINAL)
    BL.pickle.dump([relevance_effect_med, relevant_mediators_rob ], open('Relevant_mediators_rob'+ANNOTATOR+'.pckl', 'wb'))
    robust_mediators = relevant_mediators_rob.rename(index = inverse_dict).index.values.tolist()
    
    return robust_mediators


def STAGE3(compustat, regions, sector, tc, system_thinking, paris2018,paris2019,paris2020, base_dummies, MAX_ProjectionYear, ref, robust_explanators, robust_mediators, cdp_t, ref_t, cdp_firms, non_cdp_firms,  MARGINAL=False):
    MAX_ProjectionYear, paris2018,paris2019,paris2020, base_dummies = data.get_paris_for_analysis('below')
    outcome_effects = ROB.step_3_efx_robustness(compustat, regions, sector, tc, system_thinking, paris2018,paris2019,paris2020, base_dummies, MAX_ProjectionYear, ref, robust_explanators, robust_mediators, cdp_t, ref_t, cdp_firms, non_cdp_firms,  MARGINAL)
    BL.pickle.dump(outcome_effects, open('Robustness_to_sample_size'+ANNOTATOR+'_below.pckl', 'wb'))

    MAX_ProjectionYear, paris2018WELLBELOW,paris2019WELLBELOW,paris2020WELLBELOW, base_dummiesWELLBELOW = data.get_paris_for_analysis('well_below')
    outcome_effects = ROB.step_3_efx_robustness(compustat, regions, sector, tc, system_thinking, paris2018WELLBELOW,paris2019WELLBELOW,paris2020WELLBELOW, base_dummiesWELLBELOW, MAX_ProjectionYear, ref, robust_explanators, robust_mediators, cdp_t, ref_t, cdp_firms, non_cdp_firms,  MARGINAL)
    BL.pickle.dump(outcome_effects, open('Robustness_to_sample_size'+ANNOTATOR+'_well_below.pckl', 'wb'))

def RollingEstimation(system_thinking, compustat, tc, regions, sector, paris2018,paris2019,paris2020, base_dummies, MAX_ProjectionYear, ref, vars_, cdp_t, ref_t, cdp_firms, non_cdp_firms, MARGINAL=False):
    #== Temporal evolution of the effect
    for year in range(2015,2021):
        print('Year:', year)
        yearly_ST = system_thinking[system_thinking.Year <= year]
        yearly_CM = compustat[compustat.fyear <= year]
        yearly_TC = tc.copy() #[tc.cyear <= year]
        #=== First identify the correct explanators
        if os.path.exists('YearEffect/Relevant_explanators_rob'+ANNOTATOR+'_'+str(year)+'.pckl'):
            relevance_effect_expl, relevant_explanators_rob  = BL.pickle.load(open('YearEffect/Relevant_explanators_rob'+ANNOTATOR+'_'+str(year)+'.pckl', 'rb'))
        else:
            relevance_effect_expl, relevant_explanators_rob  = ROB.step_1_explanators(yearly_CM, regions, sector, yearly_TC, yearly_ST, paris2018,paris2019,paris2020, base_dummies, MAX_ProjectionYear, ref, vars_, cdp_t, ref_t, cdp_firms, non_cdp_firms)
            BL.pickle.dump([relevance_effect_expl, relevant_explanators_rob ], open('YearEffect/Relevant_explanators_rob'+ANNOTATOR+'_'+str(year)+'.pckl', 'wb'))
        robust_explanators = relevant_explanators_rob.rename(index = inverse_dict).index.values.tolist()
        #=== Then identify the correct mediators
        if os.path.exists('YearEffect/Relevant_mediators_rob'+ANNOTATOR+'_'+str(year)+'.pckl'):
            relevance_effect_med, relevant_mediators_rob  = BL.pickle.load(open('YearEffect/Relevant_mediators_rob'+ANNOTATOR+'_'+str(year)+'.pckl', 'rb'))
        else:
            relevance_effect_med, relevant_mediators_rob  = ROB.step_2_mediators( yearly_CM, regions, sector, yearly_TC,  yearly_ST, paris2018,paris2019,paris2020, base_dummies, MAX_ProjectionYear, ref, robust_explanators, vars_, cdp_t, ref_t, cdp_firms, non_cdp_firms, MARGINAL)
            BL.pickle.dump([relevance_effect_med, relevant_mediators_rob ], open('YearEffect/Relevant_mediators_rob'+ANNOTATOR+'_'+str(year)+'.pckl', 'wb'))
        robust_mediators = relevant_mediators_rob.rename(index = inverse_dict).index.values.tolist()
        
        #== Then use the relevant mediators and explanators in the models
        MAX_ProjectionYear, paris2018,paris2019,paris2020, base_dummies = data.get_paris_for_analysis('below')
        outcome_effects = ROB.step_3_efx_robustness( yearly_CM, regions, sector,  yearly_TC,  yearly_ST, paris2018,paris2019,paris2020, base_dummies, MAX_ProjectionYear, ref, robust_explanators, robust_mediators, cdp_t, ref_t, cdp_firms, non_cdp_firms,  MARGINAL, years = year)
        BL.pickle.dump(outcome_effects, open('YearEffect/Robustness_to_sample_size'+ANNOTATOR+'_'+str(year)+'_below.pckl', 'wb'))
        
        MAX_ProjectionYear, paris2018WELLBELOW,paris2019WELLBELOW,paris2020WELLBELOW, base_dummiesWELLBELOW = data.get_paris_for_analysis('well_below')
        outcome_effects = ROB.step_3_efx_robustness( yearly_CM, regions, sector,  yearly_TC,  yearly_ST, paris2018WELLBELOW, paris2019WELLBELOW, paris2020WELLBELOW, base_dummiesWELLBELOW, MAX_ProjectionYear, ref, robust_explanators, robust_mediators, cdp_t, ref_t, cdp_firms, non_cdp_firms,  MARGINAL, years = year)
        BL.pickle.dump(outcome_effects, open('YearEffect/Robustness_to_sample_size'+ANNOTATOR+'_'+str(year)+'_well_below.pckl', 'wb'))